{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import re\n",
    "import collections\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model:\n",
    "    def __init__(self,maxlen=50, \n",
    "                 vocabulary_size=20000,\n",
    "                 learning_rate=1e-3,\n",
    "                 embedding_size = 256):\n",
    "        self.output_size = embedding_size\n",
    "        self.maxlen = maxlen\n",
    "        word_embeddings = tf.Variable(\n",
    "            tf.random_uniform(\n",
    "                [vocabulary_size, embedding_size], -np.sqrt(3), np.sqrt(3)\n",
    "            )\n",
    "        )\n",
    "        self.global_step = tf.get_variable(\n",
    "            \"global_step\", shape=[], trainable=False,\n",
    "            initializer=tf.initializers.zeros())\n",
    "        self.embeddings = word_embeddings\n",
    "        self.output_layer = tf.layers.Dense(vocabulary_size, name=\"output_layer\")\n",
    "        self.output_layer.build(self.output_size)\n",
    "        \n",
    "        self.BEFORE = tf.placeholder(tf.int32,[None,maxlen])\n",
    "        self.INPUT = tf.placeholder(tf.int32,[None,maxlen])\n",
    "        self.AFTER = tf.placeholder(tf.int32,[None,maxlen])\n",
    "        self.batch_size = tf.shape(self.INPUT)[0]\n",
    "        \n",
    "        self.get_thought = self.thought(self.INPUT)\n",
    "        self.attention = tf.matmul(\n",
    "            self.get_thought, tf.transpose(self.embeddings), name = 'attention'\n",
    "        )\n",
    "        fw_logits = self.decoder(self.get_thought, self.AFTER)\n",
    "        bw_logits = self.decoder(self.get_thought, self.BEFORE)\n",
    "        self.loss = self.calculate_loss(fw_logits, self.AFTER) + self.calculate_loss(bw_logits, self.BEFORE)\n",
    "        self.optimizer = tf.train.AdamOptimizer(learning_rate).minimize(self.loss)\n",
    "        \n",
    "    def get_embedding(self, inputs):\n",
    "        return tf.nn.embedding_lookup(self.embeddings, inputs)\n",
    "        \n",
    "    def thought(self, inputs):\n",
    "        encoder_in = self.get_embedding(inputs)\n",
    "        fw_cell = tf.nn.rnn_cell.GRUCell(self.output_size)\n",
    "        bw_cell = tf.nn.rnn_cell.GRUCell(self.output_size)\n",
    "        sequence_length = tf.reduce_sum(tf.sign(inputs), axis=1)\n",
    "        rnn_output = tf.nn.bidirectional_dynamic_rnn(\n",
    "            fw_cell, bw_cell, encoder_in, sequence_length=sequence_length,\n",
    "            dtype=tf.float32)[1]\n",
    "        return sum(rnn_output)\n",
    "        \n",
    "    def decoder(self, thought, labels):\n",
    "        main = tf.strided_slice(labels, [0, 0], [self.batch_size, -1], [1, 1])\n",
    "        shifted_labels = tf.concat([tf.fill([self.batch_size, 1], 2), main], 1)\n",
    "        decoder_in = self.get_embedding(shifted_labels)\n",
    "        cell = tf.nn.rnn_cell.GRUCell(self.output_size)\n",
    "        max_seq_lengths = tf.fill([self.batch_size], self.maxlen)\n",
    "        helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "            decoder_in, max_seq_lengths, time_major = False\n",
    "        )\n",
    "        decoder = tf.contrib.seq2seq.BasicDecoder(cell, helper, thought)\n",
    "        decoder_out = tf.contrib.seq2seq.dynamic_decode(decoder)[0].rnn_output\n",
    "        return decoder_out\n",
    "        \n",
    "    def calculate_loss(self, outputs, labels):\n",
    "        mask = tf.cast(tf.sign(labels), tf.float32)\n",
    "        logits = self.output_layer(outputs)\n",
    "        return tf.contrib.seq2seq.sequence_loss(logits, labels, mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simple_textcleaning(string):\n",
    "    string = re.sub('[^A-Za-z ]+', ' ', string)\n",
    "    return re.sub(r'[ ]+', ' ', string.lower()).strip()\n",
    "\n",
    "def batch_sequence(sentences, dictionary, maxlen = 50):\n",
    "    np_array = np.zeros((len(sentences), maxlen), dtype = np.int32)\n",
    "    for no_sentence, sentence in enumerate(sentences):\n",
    "        current_no = 0\n",
    "        for no, word in enumerate(sentence.split()[: maxlen - 2]):\n",
    "            np_array[no_sentence, no] = dictionary.get(word, 1)\n",
    "            current_no = no\n",
    "        np_array[no_sentence, current_no + 1] = 3\n",
    "    return np_array\n",
    "\n",
    "def counter_words(sentences):\n",
    "    word_counter = collections.Counter()\n",
    "    word_list = []\n",
    "    num_lines, num_words = (0, 0)\n",
    "    for i in sentences:\n",
    "        words = re.findall('[\\\\w\\']+|[;:\\-\\(\\)&.,!?\"]', i)\n",
    "        word_counter.update(words)\n",
    "        word_list.extend(words)\n",
    "        num_lines += 1\n",
    "        num_words += len(words)\n",
    "    return word_counter, word_list, num_lines, num_words\n",
    "\n",
    "\n",
    "def build_dict(word_counter, vocab_size = 50000):\n",
    "    count = [['PAD', 0], ['UNK', 1], ['START', 2], ['END', 3]]\n",
    "    count.extend(word_counter.most_common(vocab_size))\n",
    "    dictionary = dict()\n",
    "    for word, _ in count:\n",
    "        dictionary[word] = len(dictionary)\n",
    "    return dictionary, {word: idx for idx, word in dictionary.items()}\n",
    "\n",
    "def split_by_dot(string):\n",
    "    string = re.sub(\n",
    "        r'(?<!\\d)\\.(?!\\d)',\n",
    "        'SPLITTT',\n",
    "        string.replace('\\n', '').replace('/', ' '),\n",
    "    )\n",
    "    string = string.split('SPLITTT')\n",
    "    return [re.sub(r'[ ]+', ' ', sentence).strip() for sentence in string]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9923"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contents = []\n",
    "with open('books/Blood_Born') as fopen:\n",
    "    contents.extend(split_by_dot(fopen.read()))\n",
    "    \n",
    "with open('books/Dark_Thirst') as fopen:\n",
    "    contents.extend(split_by_dot(fopen.read()))\n",
    "    \n",
    "len(contents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8390"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contents = [simple_textcleaning(sentence) for sentence in contents]\n",
    "contents = [sentence for sentence in contents if len(sentence) > 20]\n",
    "len(contents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['blood borna story of the supernaturalby jodie rosefirst printing september copyright jodie rose all rights reserved facebook',\n",
       " 'com jodieroseauthorjodieroseauthor',\n",
       " 'comcover design https jodieszarmachphoto',\n",
       " 'com imagescastle image photo by rob jackson on unsplashtexture image federico gutierre at unsplashdisclaimerthe book as well as the characters in this book are entirely fictional',\n",
       " 'any resemblance to actual persons living or dead is entirely coincidental']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contents[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9039"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "maxlen = 50\n",
    "vocabulary_size = len(set(' '.join(contents).split()))\n",
    "embedding_size = 256\n",
    "learning_rate = 1e-3\n",
    "batch_size = 16\n",
    "vocabulary_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.utils import shuffle\n",
    "\n",
    "stride = 1\n",
    "t_range = int((len(contents) - 3) / stride + 1)\n",
    "left, middle, right = [], [], []\n",
    "for i in range(t_range):\n",
    "    slices = contents[i * stride : i * stride + 3]\n",
    "    left.append(slices[0])\n",
    "    middle.append(slices[1])\n",
    "    right.append(slices[2])\n",
    "\n",
    "left, middle, right = shuffle(left, middle, right)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_counter, _, _, _ = counter_words(middle)\n",
    "dictionary, _ = build_dict(word_counter, vocab_size = vocabulary_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Model(vocabulary_size = len(dictionary), embedding_size = embedding_size)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 525/525 [00:53<00:00, 10.40it/s, cost=11.2]\n",
      "train minibatch loop: 100%|██████████| 525/525 [00:53<00:00, 10.44it/s, cost=8.75]\n",
      "train minibatch loop: 100%|██████████| 525/525 [00:53<00:00, 10.26it/s, cost=6.41]\n",
      "train minibatch loop: 100%|██████████| 525/525 [00:53<00:00, 10.38it/s, cost=4.48]\n",
      "train minibatch loop: 100%|██████████| 525/525 [00:53<00:00, 10.47it/s, cost=3.1] \n"
     ]
    }
   ],
   "source": [
    "for i in range(5):\n",
    "    pbar = tqdm(range(0, len(middle), batch_size), desc='train minibatch loop')\n",
    "    for p in pbar:\n",
    "        index = min(p + batch_size, len(middle))\n",
    "        batch_x = batch_sequence(\n",
    "                middle[p : index],\n",
    "                dictionary,\n",
    "                maxlen = maxlen,\n",
    "        )\n",
    "        batch_y_before = batch_sequence(\n",
    "                left[p : index],\n",
    "                dictionary,\n",
    "                maxlen = maxlen,\n",
    "        )\n",
    "        batch_y_after = batch_sequence(\n",
    "                right[p : index],\n",
    "                dictionary,\n",
    "                maxlen = maxlen,\n",
    "        )\n",
    "        loss, _ = sess.run([model.loss, model.optimizer], \n",
    "                           feed_dict = {model.BEFORE: batch_y_before,\n",
    "                                        model.INPUT: batch_x,\n",
    "                                        model.AFTER: batch_y_after,})\n",
    "        pbar.set_postfix(cost=loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('books/Driftas_Quest') as f:\n",
    "    book = f.read()\n",
    "\n",
    "book = split_by_dot(book)\n",
    "book = [simple_textcleaning(sentence) for sentence in book]\n",
    "book = [sentence for sentence in book if len(sentence) > 20][100:200]\n",
    "book_sequences = batch_sequence(book, dictionary, maxlen = maxlen)\n",
    "encoded, attention = sess.run([model.get_thought, model.attention],feed_dict={model.INPUT:book_sequences})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "so magnified was the viewer that they could see individual humans going about their business in a city street. as a stack was brought into one of the holds the precise weight was verified and a significant error from the inventory log appeared. the burley miner looked at him with eyes rolling in every direction. you are not my mother she had waited a long time for that question letting it come from the boy first. it was just a cargo ship trading amongst the alliance planets. sometime in the night someone had breached port security and placed a bio crib amongst the packages. i have always wondered how i would answer that question when you asked it. grofinglaz needs to rest. who was he really and where did he come from his first memories as a child were on a ship. we were on poraslon one of our regular stops\n"
     ]
    }
   ],
   "source": [
    "from sklearn.cluster import KMeans\n",
    "from sklearn.metrics import pairwise_distances_argmin_min\n",
    "\n",
    "n_clusters = 10\n",
    "kmeans = KMeans(n_clusters=n_clusters, random_state=0)\n",
    "kmeans = kmeans.fit(encoded)\n",
    "avg = []\n",
    "closest = []\n",
    "for j in range(n_clusters):\n",
    "    idx = np.where(kmeans.labels_ == j)[0]\n",
    "    avg.append(np.mean(idx))\n",
    "closest, _ = pairwise_distances_argmin_min(kmeans.cluster_centers_,encoded)\n",
    "ordering = sorted(range(n_clusters), key=lambda k: avg[k])\n",
    "print('. '.join([book[closest[idx]] for idx in ordering]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Important words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['do',\n",
       " 'uncontrollably',\n",
       " 'coffins',\n",
       " 'feed',\n",
       " 'retrace',\n",
       " 'busying',\n",
       " 'closing',\n",
       " 'center',\n",
       " 'seats',\n",
       " 'gilded']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indices = np.argsort(attention.mean(axis=0))[::-1]\n",
    "rev_dictionary = {v:k for k, v in dictionary.items()}\n",
    "[rev_dictionary[i] for i in indices[:10]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
